﻿/*************************************************************************************
 * 工    具：  HK_HierarchyPro
 * 
 * 描    述：  本工具用于增强 Hierarchy 面板的显示功能。
 *            通过本工具，开发者可以更直观地查看游戏对象的状态、层级和组件信息，有效地提高工作效率。
 * 
 * 版    本：  V1.4
 * 作    者：  京产肠饭
 * 
 * 创    建：  2021/11/29  V1.0
 * 
 * 更    新：  2024/09/19  V1.1
 *            1.功能新增：新增Hierarchy面板中，奇数行和偶数行，背景色的绘制功能
 *            2.功能优化：优化 Layer信息，根据Hierarchy面板宽度显示隐藏的逻辑，绘制时字体右对齐
 *            
 *            2024/09/20  V1.2
 *            1.功能新增：新增Unity常用组件图标绘制
 *            2.性能优化：组件显示功能，新增常用组件图标缓存，优化性能
 *            3.性能优化：优化DrawMarker的事件处理机制
 * 
 *            2024/12/08  V1.3
 *            1.功能优化：优化组件图标显示方式，Hierarchy窗口宽度小于180时不绘制组件图标
 *            
 *            2025/01/04  V1.4
 *            1.功能优化：组件查询使用字符串的HashSet（查询组件名），减少对具体组件类型的强依赖
 *            2.功能新增：新增自定义Monobehaviour组件绘制
 *            3.功能优化：优化 ActiveToggle 显示逻辑
 * 
 * 链    接：  https://gitee.com/JingChanChangFan/hk_-unity-tools
*************************************************************************************/

using System.Collections.Generic;
using System.Reflection;
using UnityEngine;
using UnityEditor;
using System;

namespace HKTools
{
    [InitializeOnLoad]
    public class HK_HierarchyPro
    {
        // 奇数行背景色
        static readonly Color oddLineBGColor = new Color(1f, 1f, 1f, 0.035f);
        // 偶数行背景色
        static readonly Color evenLineBGColor = new Color(0f, 0f, 0f, 0.05f);

        // 最大显示的图标数量
        const int MaxIconCount = 3;

        static Transform curOffsetTF = null;
        static int offsetCount = 0;

        // 用于记录哪些 GameObject（ID）已经被标记
        static HashSet<int> markedObjectIDs = new HashSet<int>();

        // 想要显示的 组件 列表
        static HashSet<string> showTypes = new HashSet<string>()
        {
            "Canvas",
            "Light",
            "Camera",
            "Collider",
            "AudioSource",
            "AudioListener",
        };

        // 不想显示的 组件 列表
        static HashSet<string> hideTypes = new HashSet<string>()
        {
            "HK_Button",
            "HK_Dropdown",
            "HK_Toggle",
            "UniversalAdditionalLightData",
            "UniversalAdditionalCameraData",
            "UnRegisterOnDestroyTrigger"
        };

        // 将要绘制的组件
        static List<Component> displayComponents = new List<Component>();

        // 绘制的 Component 图标大小
        static readonly int iconSize = 16;

        // 缓存常用的 Component 图标
        static Dictionary<Type, Texture2D> componentIconCache = new Dictionary<Type, Texture2D>();

        /// <summary>
        /// 用于激活或停用 Hierarchy 面板的增强显示功能
        /// </summary>
        public static void ChangeState(bool isActive)
        {
            if (isActive)
            {
                EditorApplication.hierarchyWindowItemOnGUI += OnHierarchyWindowItemOnGUI;
            }
            else
            {
                EditorApplication.hierarchyWindowItemOnGUI -= OnHierarchyWindowItemOnGUI;
            }

            EditorApplication.RepaintHierarchyWindow();
        }

        static void OnHierarchyWindowItemOnGUI(int instanceID, Rect selectionRect)
        {
            GameObject go = EditorUtility.InstanceIDToObject(instanceID) as GameObject;

            //剔除最上面的scene 物体
            if (go == null)
            {
                return;
            }

            DrawBackgroundAndOutLine(go, selectionRect);
            DrawLayerInfo(go, selectionRect);
            DrawActiveToggle(go, selectionRect);
            DrawComponentIcons(go, selectionRect);
            DrawMarker(go, selectionRect);
        }

        /// <summary>
        /// 绘制 背景色和外边框
        /// </summary>
        static void DrawBackgroundAndOutLine(GameObject go, Rect rect)
        {
            // 计算当前行的索引  
            int rowIndex = Mathf.RoundToInt(rect.y / rect.height);

            // 判断奇偶行  
            bool isOdd = rowIndex % 2 == 1;

            // 设置新的背景色  
            Color rectColor = isOdd ? oddLineBGColor : evenLineBGColor;

            // 确保绘制区域不超出Hierarchy窗口的左边界
            float drawX = Mathf.Max(rect.xMin, rect.x - 13f);

            // 绘制背景矩形  
            EditorGUI.DrawRect(new Rect(rect.x - 13f, rect.y, rect.width, rect.height), rectColor);

            Color horizontalLineColor = new Color(0.15f, 0.15f, 0.15f, 0.7f);
            Color verticalLineColor = new Color(0.6f, 0.6f, 0.6f, 0.55f);

            // 绘制上边框
            EditorGUI.DrawRect(new Rect(rect.x - 14.5f, rect.y, rect.width, 1), horizontalLineColor);
            EditorGUI.DrawRect(new Rect(rect.x - 14.5f, rect.y + rect.height, rect.width, 1), horizontalLineColor);
            EditorGUI.DrawRect(new Rect(rect.x - 14.5f, rect.y, 1, rect.height), verticalLineColor);
        }

        /// <summary>
        /// 绘制 Layer 信息
        /// </summary>
        static void DrawLayerInfo(GameObject go, Rect rect)
        {
            // 当 Hierarchy 整体宽度大约小于300时，则不显示Layer信息
            if (EditorGUIUtility.currentViewWidth < 300)
                return;

            rect.x = rect.width + rect.x - 200;
            rect.width = 110; // 设置矩形宽度，以适应文本显示  
            rect.y += 1;

            var layerInfoStyle = new GUIStyle(GUI.skin.label)
            {
                normal = { textColor = Color.gray },
                fontSize = 11,
                alignment = TextAnchor.MiddleRight
            };

            string layerName = LayerMask.LayerToName(go.layer);

            if (layerName != "Default")
            {
                GUI.Label(rect, layerName, layerInfoStyle);
            }
        }

        /// <summary>
        /// 绘制 激活状态切换
        /// </summary>
        static void DrawActiveToggle(GameObject go, Rect rect)
        {
            // 当 Hierarchy 整体宽度大约小于180时，则不显示ActiveToggle
            if (EditorGUIUtility.currentViewWidth < 180)
                return;

            rect.x = rect.width + rect.x - 15;
            rect.y += 0.4f;

            bool isOriginalActive = go.activeSelf;
            bool isCurrentActive = GUI.Toggle(rect, isOriginalActive, string.Empty);

            if (isOriginalActive != isCurrentActive)
            {
                Undo.RecordObject(go, "Toggle Active State");
                go.SetActive(isCurrentActive);
            }
        }

        /// <summary>
        /// 绘制 组件图标
        /// </summary>
        static void DrawComponentIcons(GameObject go, Rect rect)
        {
            // 当 Hierarchy 整体宽度大约小于 200 时，则不显示Component图标
            if (EditorGUIUtility.currentViewWidth < 200)
                return;

            //求出当前Hierarchy的总宽度
            rect.width += rect.x;

            //取得当前物体所有的 Component
            Component[] components = go.GetComponents<Component>();

            displayComponents.Clear();

            //判断自定义的Monobehaviour脚本，并进行显示
            for (int i = components.Length - 1; i >= 0; i--)
            {
                if (components[i].IsCustomMonoBehaviour() && !hideTypes.Contains(components[i].GetType().Name))
                {
                    displayComponents.Add(components[i]);
                }
            }

            //遍历coms中的组件类型，如果有不显示的，就移除掉
            for (int i = components.Length - 1; i >= 0; i--)
            {
                if (components[i] != null && showTypes.Contains(components[i].GetType().Name))
                {
                    if (!displayComponents.Contains(components[i]))
                        displayComponents.Add(components[i]);
                }
            }

            //如果当前物体是 需要偏移的物体  offset 为 offsetNum 否则为 0
            int offset = go.transform == curOffsetTF ? offset = offsetCount : 0;

            //如果当前物体的组件数量 <=  maxIconNum  把nowOffset 和 offsetNum 
            if (displayComponents.Count <= MaxIconCount && go.transform == curOffsetTF)
            {
                curOffsetTF = null;
                offsetCount = 0;
            }

            if (displayComponents.Count > MaxIconCount)
            {
                GUIStyle style = new GUIStyle(GUI.skin.label);
                style.fontSize = 9;
                style.alignment = TextAnchor.MiddleCenter;

                if (GUI.Button(new Rect(rect.width - (iconSize + 1.5f) * (MaxIconCount + 2), rect.y + 1, 17, 17), "•••", style))
                {
                    if (curOffsetTF != go.transform)
                    {
                        curOffsetTF = go.transform;
                        offsetCount = 0;
                    }
                    offsetCount += MaxIconCount;
                    if (offsetCount >= displayComponents.Count)
                    {
                        offsetCount = 0;
                    }
                }
            }

            for (int i = 0; i + offset < displayComponents.Count && i < MaxIconCount; i++)
            {
                Color originalColor = GUI.color;

                Component com = displayComponents[i + offset];
                PropertyInfo propertyInfo;
                try
                {
                    propertyInfo = com.GetType().GetProperty("enabled");
                }
                catch
                {
                    continue;
                }

                if ((bool)propertyInfo.GetGetMethod().Invoke(com, null))
                {
                    GUI.color = Color.white;
                }
                else
                {
                    GUI.color = Color.gray;
                }

                // 获取图标时，如果当前缓存有，则直接返回对应的 Texture2D，如果当前缓存没有，则在缓存中添加
                if (!componentIconCache.TryGetValue(com.GetType(), out Texture2D comTexture))
                {
                    //根据组件类型返回 组件的小图标
                    comTexture = AssetPreview.GetMiniThumbnail(com);
                    componentIconCache[com.GetType()] = comTexture;
                }

                if (comTexture)
                {
                    GUIStyle style = new GUIStyle();
                    style.normal.textColor = Color.white;

                    if (GUI.Button(new Rect(rect.width - (iconSize + 1) * (i + 2), rect.y + 1, iconSize, iconSize), comTexture, style))
                    {
                        bool enabled = (bool)propertyInfo.GetGetMethod().Invoke(com, null);

                        Undo.RecordObject(com, enabled ? "Disable Component" : "Enable Component");
                        propertyInfo.GetSetMethod().Invoke(com, new object[] { !enabled });
                    }
                    GUI.color = Color.white;
                }

                GUI.color = originalColor;
            }
        }

        /// <summary>
        /// 绘制 标记
        /// </summary>
        static void DrawMarker(GameObject go, Rect rect)
        {
            if (markedObjectIDs.Contains(go.GetInstanceID()))
            {
                rect.x = rect.x - 27; // 调整标记位置
                rect.width = 20;

                GUIStyle style = new GUIStyle();
                style.normal.textColor = Color.yellow;
                GUI.Label(rect, "★", style); // 使用星星符号作为标记，可以替换为其他符号
            }

            // 只有在 Hierarchy 窗口聚焦时才处理键盘事件
            if (EditorWindow.focusedWindow != null && EditorWindow.focusedWindow.titleContent.text == "Hierarchy")
            {
                Event currentEvent = Event.current;

                // 检查是否按下了 Ctrl+M
                if (currentEvent.type == EventType.KeyDown && currentEvent.keyCode == KeyCode.M && currentEvent.control)
                {
                    // 检查当前选择的 GameObject
                    foreach (GameObject obj in Selection.gameObjects)
                    {
                        int id = obj.GetInstanceID();

                        if (markedObjectIDs.Contains(id))
                        {
                            markedObjectIDs.Remove(id);
                        }
                        else
                        {
                            markedObjectIDs.Add(id);
                        }
                    }

                    // 重新绘制 Hierarchy 窗口
                    EditorApplication.RepaintHierarchyWindow();

                    // 使用事件，防止后续处理
                    Event.current.Use();
                }
            }
        }
    }

    public static class ComponentExtensions
    {
        /// <summary>  
        /// 判断组件是否是自定义的 MonoBehaviour 脚本。  
        /// </summary>  
        public static bool IsCustomMonoBehaviour(this Component component)
        {
            if (component == null)
                return false;

            // 检查是否继承自 MonoBehaviour  
            if (!(component is MonoBehaviour))
                return false;

            // 获取组件类型的程序集名称  
            string assemblyName = component.GetType().Assembly.GetName().Name;

            // Unity 内置程序集名称通常以 "UnityEngine" 或 "UnityEditor" 开头  
            if (assemblyName.StartsWith("UnityEngine") || assemblyName.StartsWith("UnityEditor"))
                return false;

            // 其他情况视为自定义脚本  
            return true;
        }
    }
}